#
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
#

import logging
import sys
import unittest
import math
import numpy as np

sys.path.append('../')
import torch
import random
import sys
import numpy as np
import math

np.random.seed(1)
random.seed(1)
MAX_SEQ_LEN = 1024


class PagedAttentionDataGenerator():
    def compare_output_data(self, out, golden, ratios):
        error_count = 0
        strict_error_count = 0
        fp16_min_normal = 1.0 / (1 << 14)
        golden = golden.to(torch.float32)
        out = out.to(torch.float32)
        len = out.shape[0] * out.shape[1] * out.shape[2]
        diff = torch.abs(golden - out)
        max_diff = diff.max().item()
        limit_error = torch.maximum(torch.abs(golden * ratios[0]), torch.tensor(ratios[1]))
        strict_limit_error = torch.maximum(torch.abs(golden * ratios[2]), torch.tensor(ratios[3]))
        error_count = torch.gt(diff, limit_error).sum().item()
        strict_error_count = torch.gt(diff, strict_limit_error).sum().item()
        print(f"maxDiff {max_diff}")
        print("1/1000 Accuracy is %f", 1 - float(error_count) / len)
        print("5/1000 Accuracy is %f", 1 - float(strict_error_count) / len)
        # 旧精度标准：双千分之五
        if self.data_type == torch.bfloat16 or self.is_int8_flag:
            print("accuracy is correct: %r", (float(strict_error_count) / len) <= ratios[2])
        else:
            print("accuracy is correct: %r", (float(strict_error_count) / len) <= ratios[0])
        # 新精度标准 参考精度标准v0.3浮点计算单标杆
        # 计算次数 两个matmul + 一个softmax
        calc_times = out.shape[2] * self.max_context_len + 4
        if self.data_type == torch.bfloat16:
            if calc_times < 2048:
                error = 2 ** (-7)
            else:
                error = 2 ** (-6)
            error_threshold = torch.clamp(torch.abs(golden), min=1) * error
            return (diff <= error_threshold).all()
        else:
            if calc_times < 2048:
                error = 2 ** (-8)
            else:
                error = 2 ** (-7)
            error_threshold = torch.clamp(torch.abs(golden), min=1) * error
            return (diff <= error_threshold).all()

    def group_mm_torch(self, heads, group_num, A, B, is_k):
        group_head = heads // group_num
        score = None
        for i in range(group_num):
            if self.is_int8_flag:
                int8_B = B[i: (i + 1), :, :, ]
                head_dim = int8_B.shape[2]
                int32_B = torch.matmul(torch.eye(int8_B.shape[1]).to(torch.int32), int8_B.to(torch.int32)).to(
                    torch.int32)
                if is_k:
                    if self.has_bias:
                        int32_B = int32_B + self.offset1[i * head_dim:(i + 1) * head_dim]
                    fp32_B = int32_B.to(torch.float32) * self.de_scale1_fp32[i * head_dim:(i + 1) * head_dim]
                    fp32_B = torch.permute(fp32_B, (0, 2, 1))
                else:
                    if self.has_bias:
                        int32_B = int32_B + self.offset2[i * head_dim:(i + 1) * head_dim]
                    fp32_B = int32_B.to(torch.float32) * self.de_scale2_fp32[i * head_dim:(i + 1) * head_dim]
                group_score = torch.matmul(A[i * group_head: (i + 1) * group_head, :, :].to(torch.float32),
                                           fp32_B).to(torch.float16)
            else:
                group_score = torch.matmul(A[i * group_head: (i + 1) * group_head, :, :].to(torch.float32),
                                           B[i:(i + 1), :, :].to(torch.float32))
            if score is None:
                score = group_score
            else:
                score = torch.cat((score, group_score), 0)
        return score

    def process_deq_scale(self, deq_scale) -> np.ndarray:
        new_deq_scale = np.frombuffer(deq_scale.tobytes(), dtype=np.uint32)
        return new_deq_scale.astype(np.uint64)

    def ref_masked_attention(self,
                             query,  # (1, num_heads, head_size)
                             key,  # (context_len, kv_heads, head_size)
                             value,
                             scale: float,
                             alibi_bias,
                             mask_data_type=torch.bfloat16
                             ):
        # Q * K.T
        query = query
        query = torch.permute(query, (1, 0, 2))
        if not self.is_int8_flag:
            key = torch.permute(key, (1, 2, 0))  # 0 1 2
        else:
            key = torch.permute(key, (1, 0, 2))
        sim = self.group_mm_torch(query.shape[0], key.shape[0], query, key, 1).to(
            mask_data_type)  # (head_num, q_seqlen, k_seqlen)
        if alibi_bias is None:
            sim = sim.to(torch.float16) * np.float16(scale)
        else:
            if (mask_data_type == torch.bfloat16):
                sim = sim.to(torch.float32) * scale
                sim = sim + alibi_bias.to(torch.float32)
            else:
                sim = sim.to(torch.float16) * np.float16(scale)
                sim = sim + alibi_bias.to(torch.float16)
        sim = sim.numpy()
        # softmax
        row_max = np.max(sim, axis=-1, keepdims=True)
        sim -= row_max
        sim = np.exp(sim)
        row_sum = np.sum(sim, axis=-1, keepdims=True)
        p = sim / row_sum
        p = torch.from_numpy(p).to(mask_data_type)
        # P * V
        value = torch.permute(value, (1, 0, 2))
        out = self.group_mm_torch(query.shape[0], key.shape[0], p, value, 0)
        out = torch.permute(out, (1, 0, 2))
        return out

    def ref_single_query_cached_kv_attention(self,
                                             output,
                                             query,
                                             key_cache,  # (num_blocks, block_size, num_heads, head_size)
                                             value_cache,  # (num_blocks, block_size, num_heads, head_size)
                                             block_tables,
                                             context_lens,
                                             mask,
                                             mask_dim=4,
                                             mask_data_type=torch.bfloat16
                                             ) -> None:
        mask_index_coff = 1
        head_size = key_cache.shape[3]
        head_size_v = value_cache.shape[3]
        if self.compressHead:
            query = query.view(self.num_tokens * self.kv_heads, self.num_heads // self.kv_heads, self.head_size)
            output = output.view(self.num_tokens * self.kv_heads, self.num_heads // self.kv_heads, head_size_v)
            if mask_dim == 4:
                mask_shape = mask.shape
                mask = mask.view(mask_shape[0] * self.kv_heads, self.num_heads // self.kv_heads, 1,
                                 self.max_context_len)
            else:
                mask_index_coff = self.kv_heads
        num_heads = query.shape[1]
        kv_heads = value_cache.shape[2]
        block_size = value_cache.shape[1]

        num_input_tokens = query.shape[0]
        index = 0
        for i in range(len(context_lens)):
            block_table = block_tables[i]
            context_len = int(context_lens[i])
            if context_len == 0:
                continue

            q = query[index].view(1, num_heads, head_size)
            keys = []
            values = []
            for j in range(context_len):
                block_number = int(block_table[j // block_size])
                block_offset = j % block_size

                k = key_cache[block_number, block_offset, :, :]
                k = k.reshape(kv_heads, head_size)
                keys.append(k)

                v = value_cache[block_number, block_offset, :, :]
                v = v.reshape(kv_heads, head_size_v)
                values.append(v)
            keys = torch.stack(keys, axis=0)
            values = torch.stack(values, axis=0)
            scale = np.float32(1.0 / (head_size ** 0.5))
            if mask_dim == 4:
                logging.info(f"query.shape: {q.shape}, {q.dtype}, keys.shape: {keys.shape}, "
                             f"context_len: {context_len}, keyblocknum: {(context_len + block_size - 1) // block_size}, "
                             f"tail: {context_len % block_size}, alibi_bias.shape: {mask[i].shape}")
                out = self.ref_masked_attention(q, keys, values, scale, mask[i, :, :, :context_len], mask_data_type)
                out = out.reshape(num_heads, head_size_v)
            elif mask_dim == 3:
                logging.info(f"query.shape: {q.shape}, {q.dtype}, keys.shape: {keys.shape}, "
                             f"context_len: {context_len}, keyblocknum: {(context_len + block_size - 1) // block_size}, "
                             f"tail: {context_len % block_size}, alibi_bias.shape: {mask[i // mask_index_coff].shape}")
                out = self.ref_masked_attention(q, keys, values, scale, mask[i // mask_index_coff, :, :context_len],
                                                mask_data_type)
                out = out.reshape(num_heads, head_size_v)
            else:
                out = self.ref_masked_attention(q, keys, values, scale, mask, mask_data_type)
                out = out.reshape(num_heads, head_size_v)

            output[index] = out
            index = index + 1

    def calc_data(self, num_tokens, num_heads, kv_heads, head_size, block_size, num_blocks, k_seqlen, \
                  dtype, mask_dim=4, mask_data_type=torch.bfloat16, \
                  dynamic_batch=False, dynamic_seqlen=None, is_int8_flag=False, has_bias=False,
                  compressHead=False):
        self.num_heads = num_heads
        self.kv_heads = kv_heads
        self.num_tokens = num_tokens
        self.compressHead = compressHead
        self.head_size = head_size
        head_size_v = head_size
        logging.info(
            f'input info: {num_tokens}, {num_heads}, {kv_heads}, {head_size}, {block_size}, {num_blocks}, {k_seqlen}, {dtype}')

        query = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(num_tokens, num_heads, head_size))).to(
            mask_data_type)
        # (num_blocks, block_size, num_heads, head_size)
        kv_range = 1.0
        kv_type = mask_data_type
        if is_int8_flag:
            kv_range = 4.0
            kv_type = torch.int8
        if not compressHead:
            key_cache = torch.from_numpy(
                np.random.uniform(-kv_range, kv_range, size=(num_blocks, block_size, kv_heads, head_size))).to(kv_type)
            # (num_blocks, block_size, num_heads, head_size)
            value_cache = torch.from_numpy(
                np.random.uniform(-kv_range, kv_range, size=(num_blocks, block_size, kv_heads, head_size_v))).to(
                kv_type)
        else:
            key_cache = torch.from_numpy(
                np.random.uniform(-kv_range, kv_range, size=(num_blocks * kv_heads, block_size, 1, head_size))).to(
                kv_type)
            # (num_blocks, block_size, num_heads, head_size)
            value_cache = torch.from_numpy(
                np.random.uniform(-kv_range, kv_range, size=(num_blocks * kv_heads, block_size, 1, head_size_v))).to(
                kv_type)
        self.data_type = dtype

        if dynamic_batch:
            context_lens = dynamic_seqlen
        else:
            context_lens = [k_seqlen] * num_tokens
        max_context_len = max(context_lens)
        self.max_context_len = max_context_len
        batch = len(context_lens)

        # alibi mask
        if mask_dim == 4:
            alibi_slopes = np.random.random(num_heads).astype(np.float16)
            mask = np.zeros((batch, num_heads, 1, max_context_len), dtype=np.float16)
            for i, context_len in enumerate(context_lens):
                if context_len == 0:
                    continue
                position_ids = np.arange(context_len).astype(np.int32)
                alibi_bias = (position_ids - context_len + 1).astype(np.float16)
                alibi_bias = alibi_slopes.reshape(-1, 1, 1) * alibi_bias.reshape(1, 1, -1)  # (head_num, 1, context)
                mask[i, :, :, :context_len] = alibi_bias
            logging.info(f"mask.shape = {mask.shape}")
            mask = torch.from_numpy(mask).to(mask_data_type)
        # normal mask
        elif mask_dim == 3:
            mask = np.zeros((batch, 1, max_context_len), dtype=np.float16)
            for i in range(batch):
                mask[i, :, :i] = -10000
            logging.info(f"norm_mask.shape = {mask.shape}")
            mask = torch.from_numpy(mask).to(mask_data_type)
        else:  # no mask
            mask = None

        if compressHead:
            context_lens = [val for val in context_lens for _ in range(kv_heads)]
        batch = len(context_lens)
        max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
        block_tables = []  # （num_tokens, max_num_blocks_per_seq）
        for _ in range(batch):
            block_table = [
                random.randint(0, num_blocks - 1) for _ in range(max_num_blocks_per_seq)
            ]
            block_tables.append(block_table)

        self.is_int8_flag = is_int8_flag
        if is_int8_flag:
            de_scale1_fp32 = np.random.randint(-1, 2, size=(kv_heads * head_size)).astype(np.float32)
            de_scale1_int64 = self.process_deq_scale(de_scale1_fp32)

            de_scale2_fp32 = np.random.randint(-1, 2, size=(kv_heads * head_size)).astype(np.float32)
            de_scale2_int64 = self.process_deq_scale(de_scale2_fp32)

            offset1 = np.random.randint(-20, 20, size=(kv_heads * head_size)).astype(np.int32)

            offset2 = np.random.randint(-20, 20, size=(kv_heads * head_size)).astype(np.int32)

            self.de_scale1_int64 = torch.tensor(list(de_scale1_int64), dtype=torch.int64)
            self.de_scale2_int64 = torch.tensor(list(de_scale2_int64), dtype=torch.int64)
            self.de_scale1_fp32 = torch.from_numpy(de_scale1_fp32)
            self.de_scale2_fp32 = torch.from_numpy(de_scale2_fp32)
            self.offset1 = torch.from_numpy(offset1)
            self.offset2 = torch.from_numpy(offset2)
            self.has_bias = has_bias

        ref_output = torch.zeros((num_tokens, num_heads, head_size_v))
        ref_output[:] = 0.1
        self.ref_single_query_cached_kv_attention(
            ref_output,
            query,
            key_cache,
            value_cache,
            block_tables,
            context_lens,
            mask,
            mask_dim,
            mask_data_type
        )

        self.q = query
        self.key_cache = key_cache
        self.value_cache = value_cache
        self.block_tables = np.array(block_tables).astype(np.int32)
        self.contex_lens = np.array(context_lens).astype(np.int32)
        self.alib_mask = mask
        self.golden_out = ref_output

    def test_pa_fp16_case_norm_mask_multi_batch_compress_head(self):

        compressHead = True
        num_tokens = 1
        num_heads = 8
        kv_heads = 8
        block_size = 128
        head_size = 128
        num_blocks = 64
        seqlen = 256
        k_seqlen = seqlen if compressHead else seqlen * 2
        tor = 1.0 / (head_size ** 0.5)
        dtype = torch.float16
        mask_dim = 3

        self.calc_data(num_tokens, num_heads, kv_heads, head_size, block_size, num_blocks, k_seqlen, dtype, mask_dim,
                       torch.float16, compressHead=compressHead)